#include <stdio.h>
#include <unistd.h>
#include <errno.h>
#include <string.h>

#include <sys/types.h>
#include <sys/time.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <netinet/ip.h>

#include <linux/icmp.h>
#include <linux/errqueue.h>

#include <stdint.h>
#include <unistd.h>

#include <arpa/inet.h>

#include <poll.h>

// ================================================================================

typedef struct icmphdr icmphdr_t;

typedef struct sockaddr_in sockaddr_in_t;
typedef struct sockaddr    sockaddr_t;

// ================================================================================

static uint16_t ip_checksum(uint16_t *data, int len) {
  register int nleft = len;
  register uint16_t answer;
  register uint32_t sum = 0;

  /*
   *  Our algorithm is simple, using a 32 bit accumulator (sum),
   *  we add sequential 16 bit words to it, and at the end, fold
   *  back all the carry bits from the top 16 bits into the lower
   *  16 bits.
   */
  while (nleft > 1)  {
    sum += *data++;
    nleft -= 2;
  }
  
  /*
   * add back carry outs from top 16 bits to low 16 bits
   */
  sum = (sum >> 16) + (sum & 0xffff);	/* add hi 16 to low 16 */
  sum += (sum >> 16);			/* add carry */
  answer = ~sum;				/* truncate to 16 bits */
  return (answer);
}


static void usage() {
  printf("Usage: traceroute <Host>\n");
}

int main(int argc, char *argv[]) {
  int probe_sock, icmp_sock;
  sockaddr_in_t dest, host;
  unsigned int ttl;
  struct timeval timeout;
  uint16_t seqn;
  int id;

  static char buf[64];
  static char recv_buf[1024];
  static char control_buf[4096];
  
  if (argc < 2) {
    usage();
    return 1;
  }

  if (inet_aton(argv[1], &dest.sin_addr) < 0) {
    usage();
    return 1;
  }
  dest.sin_family = AF_INET;
  dest.sin_port = 0;
  
  host.sin_family = AF_INET;
  host.sin_port = 0;
  host.sin_addr.s_addr = INADDR_ANY;
  icmp_sock = socket(AF_INET, SOCK_RAW, IPPROTO_ICMP);
  if (icmp_sock == -1) {
    perror("Failed to create an ICMP socket");
    return 1;
  }
  if (bind(icmp_sock, (sockaddr_t*)&host, sizeof(host)) == -1) {
    perror("Failed to bind to the ICMP protocol");
    goto close_icmp;
  }
  if (connect(icmp_sock, (sockaddr_t*)&dest, sizeof(dest)) == -1) {
    perror("Failed to switch the ICMP socket into the connected state");
    goto close_icmp;
  }

  memset(buf, 'A', sizeof(buf));

  /* 
   *  Probably we don't have to setup the ICMP filter.. 
   */
  struct icmp_filter flt;
  flt.data = ~((1<<ICMP_SOURCE_QUENCH)|
               (1<<ICMP_DEST_UNREACH)|
               (1<<ICMP_TIME_EXCEEDED)|
               (1<<ICMP_PARAMETERPROB)|
               (1<<ICMP_REDIRECT)|
               (1<<ICMP_ECHOREPLY));
  if (setsockopt(icmp_sock, SOL_RAW, ICMP_FILTER, &flt, sizeof(flt))) {
    perror("Failed to set an ICMP filter");
    goto close_icmp;
  }

  
  timeout.tv_sec = 0;
  timeout.tv_usec = 100000;
  setsockopt(icmp_sock, SOL_SOCKET, SO_RCVTIMEO, &timeout, sizeof(timeout));
  setsockopt(icmp_sock, SOL_SOCKET, SO_SNDTIMEO, &timeout, sizeof(timeout));

  /*
   *  I don't know whether we need it --- the original ping does it.
   *  Maybe someone wants to find out :)
   */
  int val = 1;
  setsockopt(icmp_sock, SOL_IP, IP_RECVERR, &val, sizeof(val));
  
  val = 324;
  setsockopt(icmp_sock, SOL_SOCKET, SO_SNDBUF, &val, sizeof(val));
  val = 131072;  
  setsockopt(icmp_sock, SOL_SOCKET, SO_RCVBUF, &val, sizeof(val));
  //val = 1;
  //setsockopt(icmp_sock, SOL_SOCKET, SO_TIMESTAMP, &val, sizeof(val));
  
  ttl = 1;
  seqn = 0;
  id = getpid() & 0xFFFF;
  for (;;) {
    struct icmphdr* icmp_msg = (struct icmphdr*)buf;
    struct msghdr msg;
    struct iovec  iov;

  next_probe:

    seqn++;        
        
    setsockopt(icmp_sock, IPPROTO_IP, IP_TTL, &ttl, sizeof(ttl));

    icmp_msg = (struct icmphdr*)buf;
    icmp_msg->type = ICMP_ECHO;
    icmp_msg->code = 0;
    icmp_msg->un.echo.id = id;
    icmp_msg->un.echo.sequence = htons(seqn);
    icmp_msg->checksum = 0;
    icmp_msg->checksum = ip_checksum((uint16_t*)buf, sizeof(buf));
    send(icmp_sock, buf, sizeof(buf), 0);
    
    usleep(100000);


    iov.iov_base = &recv_buf;
    iov.iov_len = sizeof(recv_buf);

    memset(recv_buf, 0, sizeof(recv_buf));
    memset(control_buf, 0, sizeof(control_buf));

    msg.msg_name = &host;
    msg.msg_namelen = sizeof(host);
    msg.msg_iov = &iov;
    msg.msg_iovlen = 1;
    msg.msg_control = control_buf;
    msg.msg_controllen = sizeof(control_buf);
    msg.msg_flags = 0;

    for (;;) {
      struct pollfd pfd;
      int res;

      pfd.fd = icmp_sock;
      pfd.events = POLLIN | POLLERR;
      res = poll(&pfd, 1, 1000);

      if (res > 0) { 
        if (pfd.revents & POLLERR) {   

          /*
           *  How could I know?..
           */ 
          if (recvmsg(icmp_sock, &msg, MSG_ERRQUEUE) != -1) {
            struct cmsghdr* cmsg;

            /* 
             *  A black magic from the kernel space :) 
             *  The <linux/errqueue.h> doesn't want to explain this.
             *  The original traceroute doesn't want to explain this either...
             */
            for (cmsg = CMSG_FIRSTHDR(&msg);
                 cmsg;
                 cmsg = CMSG_NXTHDR(&msg, cmsg)) {
              if (cmsg->cmsg_level == SOL_IP && cmsg->cmsg_type == IP_RECVERR) {
                sockaddr_in_t* addrp = (sockaddr_in_t*)SO_EE_OFFENDER((struct sock_extended_err*)CMSG_DATA(cmsg));
                struct icmphdr* reply = (struct icmphdr*)recv_buf;

                if (reply->un.echo.id == id) {
                  printf("%d: %s\n", ttl, inet_ntoa(addrp->sin_addr));
                  ttl++;
                  id++;
                }
              }
            } 
          } else {
            perror("Error receiving a error message from the error queue");
            break;
          }
        } else if (pfd.revents & POLLIN) {
          // Received an ECHO ICMP Reply?
          printf("Tracerouting finished.\n");
          goto close_icmp;
        }
      } else if (res == 0) {
        break;          
      } else {
        perror("Error recieving a message from the ICMP socket");
        usleep(100000);
        break;
      }
    }
  }

  
 close_icmp:
  close(icmp_sock);

  return 0;
}
      
      
